'''
Created on Jul 8, 2009

@author: morales
'''

import summaryStats as ss
from math import sqrt

class TajimasD:
    
    def calculateNumerator(self, SFS):
        pi = ss.heterozygosity(SFS)
        s = ss.segregatingSites(SFS)
        
        n = SFS.shape[0]
        a = self.a_n(n)
        num = 1.0 * (pi - (s / a))
        
        return num

    def calculate(self, SFS):

        num = self.calculateNumerator(SFS)
        
        s = ss.segregatingSites(SFS)
        n = SFS.shape[0]
        
        den = self.variance( n, s )
        D = 1.0 * num / sqrt(den)
        
        return D

    def a_n(self, n):
        sum = 0.0
    
        for i in range( 1, n ):
            sum += 1.0 / i
    
        return sum
    
    def b_n(self, n):
        sum = 0.0
    
        for i in range( 1, n ):
            sum += 1.0 / (i*i)
    
        return sum
    
    def variance( self, n, s ):

        e1n = self.e1_n( n )
        e2n = self.e2_n( n )

        v = 1.0 * (e1n * s) + (e2n * s * (s-1.0))
        return v
   
    def e1_n( self, n ):
      
        a = self.a_n( n )
      
        e = 1.0 * (n+1.0) / ( 3.0 * (n-1.0) )
        e -= 1.0 / a
        e /= a
      
        return e
   
    def e2_n( self, n ):
      
        a = self.a_n( n )
        b = self.b_n( n )
      
        e = 2.0 * (n*(n+1.0) + 3.0) / (9.0 * n * (n-1.0))
        e -= 1.0 * (n+2.0) / (n*a)
        e += 1.0 * b / (a*a)
        e /= 1.0 * a*a + b
      
        return e
